import argparse
import json
import math
import os
import re
import pdb

import matplotlib.pyplot as plt
import numpy as np

GPU_JSONL_FILE = 'gpu.jsonl'
NPU_LOG_FILE = 'npu.log'
FIG_PATH = 'loss_lr_plot.png'


def get_npu_data_from_log(npu_log_file):
    npu_data = []
    with open(npu_log_file, 'r') as f:
        for line in f:
            if 'iteration' in line and 'lm loss' in line:
                match = re.search(r'iteration\s+(\d+)/\s+.*?lm loss:\s+([\d.E+-]+)', line)
                if match:
                    iteration = int(match.group(1))
                    lm_loss = float(match.group(2))
                    lr_match = re.search(r'learning rate:\s+([\d.E+-]+)', line)
                    lr = float(lr_match.group(1)) if lr_match else None
                    npu_data.append({
                        'current_steps': iteration,
                        'loss': lm_loss,
                        'lr': lr
                    })
    return npu_data


def get_gpu_data_from_jsonl(gpu_jsonl_file):
    gpu_data = []
    with open(gpu_jsonl_file, 'r') as f:
        for line in f:
            gpu_data.append(json.loads(line))
    return gpu_data


def draw_loss_lr(gpu_jsonl_file, npu_log_file, output_path=FIG_PATH):
    print(f'====get gpu_data from {gpu_jsonl_file}===')
    gpu_data = get_gpu_data_from_jsonl(gpu_jsonl_file)
    print(f'gpu_data:{gpu_data}')

    print(f'====get npu_data from {npu_log_file}===')
    npu_data = get_npu_data_from_log(npu_log_file)
    print(f'npu_data:{npu_data}')

    # extract data
    steps_gpu = [d['current_steps'] for d in gpu_data if d.get('loss')]
    # pdb.set_trace()
    loss_gpu = [d['loss'] for d in gpu_data if d.get('loss')]
    lr_gpu = [d['lr'] for d in gpu_data if d.get('lr')]
    steps_npu = [d['current_steps'] for d in npu_data]
    loss_npu = [d['loss'] for d in npu_data]
    lr_npu = [d['lr'] for d in npu_data]
    # get common steps
    common_steps = list(set(steps_gpu) & set(steps_npu))
    common_steps.sort()

    # get common loss和lr
    # pdb.set_trace()
    loss_gpu_common = [loss_gpu[steps_gpu.index(s)] for s in common_steps]
    loss_npu_common = [loss_npu[steps_npu.index(s)] for s in common_steps]
    #print(loss_gpu_common)
    #print(loss_npu_common)
    lr_gpu_common = [lr_gpu[steps_gpu.index(s)] for s in common_steps]
    lr_npu_common = [lr_npu[steps_npu.index(s)] for s in common_steps]

    # get loss_abs_error
    loss_error = [l2 - l1 for l1, l2 in zip(loss_gpu_common, loss_npu_common)]
    lr_error = [l2 - l1 for l1, l2 in zip(lr_gpu_common, lr_npu_common)]
    loss_abs_error = [abs(l2 - l1) for l1, l2 in zip(loss_gpu_common, loss_npu_common)]
    lr_abs_error = [abs(l2 - l1) for l1, l2 in zip(lr_gpu_common, lr_npu_common)]
    print(f'loss_abs_error:{loss_error}')
    print(f'lr_abs_error:{lr_error}')

    one_epoch = math.ceil(len(loss_error) // 5)
    loss_abs_error_one_epoch = [abs(l1 - l2) for l1, l2 in zip(loss_gpu_common[-one_epoch:], loss_npu_common[-one_epoch:])]
    #lr_abs_error = [abs(l1 - l2) for l1, l2 in zip(lr_gpu_common, lr_npu_common)]
    print(f'loss_abs_error_one_epoch({len(loss_abs_error_one_epoch)}):{loss_abs_error_one_epoch}')
    print(f'loss_abs_error_one_epoch_ave:{np.mean(loss_abs_error_one_epoch)}')
    #print(f'lr_abs_error:{lr_abs_error}')

    # creat fig
    plt.figure(figsize=(15, 10))
    gpu_label = os.path.basename(gpu_jsonl_file)
    npu_label = os.path.basename(npu_log_file)
    # draw loss
    plt.subplot(2, 2, 1)
    plt.plot(common_steps, loss_gpu_common, label=gpu_label, color='blue')
    plt.plot(common_steps, loss_npu_common, label=npu_label, color='red')
    plt.xlabel('Steps')
    plt.ylabel('Loss')
    plt.title('Loss Comparison')
    plt.legend()
    plt.grid(True)
    # draw loss error
    plt.subplot(2, 2, 2)
    plt.plot(common_steps, loss_error, label='Loss Absolute Error', color='green')
    plt.xlabel('Steps')
    plt.ylabel('Absolute Error')
    mean_loss_abs_error = np.mean(loss_abs_error)
    plt.title(f'Loss Absolute Error={mean_loss_abs_error:.2e}')
    # plt.text(mean_loss_error, mean_loss_error, f'{mean_loss_error:.2e}', verticalalignment='bottom', color='blue')
    plt.legend()
    plt.grid(True)

    # draw lr
    plt.subplot(2, 2, 3)
    plt.plot(common_steps, lr_gpu_common, label=gpu_label, color='blue')
    plt.plot(common_steps, lr_npu_common, label=npu_label, color='red')
    plt.xlabel('Steps')
    plt.ylabel('Learning Rate')
    plt.title('Learning Rate Comparison')
    plt.legend()
    plt.grid(True)

    # draw lr error
    plt.subplot(2, 2, 4)
    plt.plot(common_steps, lr_error, label='Learning Rate Absolute Error', color='green')
    plt.xlabel('Steps')
    plt.ylabel('Absolute Error')
    mean_lr_abs_error = np.mean(lr_abs_error)
    plt.title(f'Learning Rate Absolute Error={mean_lr_abs_error:.2e}')

    # plt.text(mean_lr_error, mean_lr_error, f'{mean_lr_error:.2e}', verticalalignment='bottom', color='blue')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    # plt.show()
    plt.savefig(output_path, dpi=100, bbox_inches='tight')
    print(f"\n\n=== Plot saved: {output_path} ===")
    plt.close()


def main():
    parser = argparse.ArgumentParser(
        description='loss,lr diff between npu_log and gpu_jsonl ')
    parser.add_argument('npu_log', help='npu log file')
    parser.add_argument('gpu_json', help='gpu jsonl file')
    parser.add_argument('-o', '--fig_file', default=FIG_PATH,
                        help='Output directory (default: current)')
    args = parser.parse_args()

    draw_loss_lr(args.npu_log, args.gpu_json, output_path=args.fig_file)


if __name__ == "__main__":
    main()
